#include "math.h"
#include <iostream>
#include "mylib.h"

using namespace mylib;

VecDoub f(VecDoub y, double t)
{
    int n = y.size();
    VecDoub dy(n, 0.0);

    dy[0] = y[2];
    dy[1] = y[3];
    //    dy[2] = 0;
    //    dy[3] = -9.8;
    dy[2] = -0.004 * y[2] * y[2];
    dy[3] = -9.8 + 0.004 * y[3] * y[3];

    return dy;
}

void DSolveEular(VecDoub y0, double tau, int n, VecDoub y[], VecDoub (*f)(VecDoub, double))
{
    y[0] = y0;

    for (int i = 0; i < n - 1; i++)
    {
        VecDoub tmp = (f(y[i], i * tau) * tau);
        y[i + 1] = y[i] + tmp;
        y[i + 1].show(4);
    }
}

void DSolveRK(VecDoub y0, double tau, int n, VecDoub y[], VecDoub (*f)(VecDoub, double))
{
    y[0] = y0;

    for (int i = 0; i < n - 1; i++)
    {
        VecDoub k1 = (f(y[i], i * tau) * tau) * 0.5;
        VecDoub k2 = (f(y[i] + k1, i * tau + 0.5 * tau) * tau) * 0.5;
        VecDoub k3 = (f(y[i] + k2, i * tau + 0.5 * tau) * tau);
        VecDoub k4 = (f(y[i] + k3, (i + 1) * tau) * tau);
        k1 = k1 * (1.0 / 3.0);
        k2 = k2 * (2.0 / 3.0);
        k3 = k3 * (1.0 / 3.0);
        k4 = k4 * (1.0 / 6.0);
        y[i + 1] = y[i] + k1 + k2 + k3 + k4;
        y[i + 1].show(4);
    }
}

int main(int argc, char const *argv[])
{
    Doub tmp[] = {0, 100, 50, 0};
    VecDoub y0(4, tmp);
    int n = 100;
    double tau = 0.1;
    VecDoub y[n];

    std::cout<<"#====== Eular======"<<std::endl;
    DSolveEular(y0, tau, n, y, f);
    std::cout<<"#====== RK======"<<std::endl;
    DSolveRK(y0, tau, n, y, f);
    return 0;
}
